from typing import Dict, List, Optional

from enum import Enum
from ..format_utils import format_str_list_to_a_string
from .EntityFormatter import EntityFormatter
from ...core.datatypes import Entity


class Language(Enum):
    EN = "en"


class AnswerFormatter:
    def __init__(
        self,
        sep: str,
        use_special_token_for_label: bool = False,
        label2token: Dict[str, str] = None,
        template_order: int = 0,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.use_special_token_for_label = use_special_token_for_label
        self.label2token = label2token
        self.sep = sep
        self.template_order = template_order

    available_languages = [Language.EN.value]
    patterns = {
        Language.EN.value: "{0} {1} {2}",  # sub rel_type obj
    }

    def extract_triples(self, relations: List[Dict[str, List[str]]]):
        """Extract triples from relations

        Args:
            relations (List[Dict[str, List[str]]]): 关系列表， 其格式为
                [
                    {
                        "head": {"name": "Lionel Messi", "type": "footballer"},
                        "type": "work for",
                        "tail": {"name": "Barcelona", "type": "football club"}
                    }
                ]

        Returns:
            List[Tuple[str, str, str]]: 返回三元组列表，其格式为 [(head, rel_type, tail)]
        """
        triples = []
        for r in relations:
            head = r["head"]
            tail = r["tail"]
            rel_type = (
                self.label2token.get(r["type"].lower())
                if self.use_special_token_for_label
                else r["type"].lower()
            )

            triples.append((head["name"], rel_type, tail["name"]))

        return triples

    def from_relations_to_triples(cls, relations: Optional[List[Dict[str, List[str]]]]):
        """输出模板为三元组形式，即 (head, rel_type, tail)"""
        triples = cls.extract_triples(relations=relations)

        return cls.sep.join([f"({','.join(triple)})" for triple in triples])

    def from_relations_to_subject_verb_objext_structure(
        cls, relations: Optional[List[Dict[str, List[str]]]], language: str = "en"
    ) -> Optional[str]:
        """输出模板为主谓宾结构，即 head rel_type tail"""

        if language not in cls.available_languages:
            raise ValueError(
                f"Expected language to be one of {cls.available_languages}"
            )

        triples = cls.extract_triples(relations=relations)

        answers = [
            cls.patterns[language].format(head, rel_type, tail)
            for head, rel_type, tail in triples
        ]

        return format_str_list_to_a_string(answers, cls.sep)

    def format_relations_to_answer_template(
        cls,
        relations: Optional[List[Dict[str, List[str]]]],
    ) -> Optional[str]:
        """输出模板为答案模板，即 the relation between {head} and {subject} is {relation}"""

        if relations is None:
            return None

        triples = cls.extract_triples(relations=relations)
        answers = [
            f"The relation between {head['name']} and {tail['name']} is {rel_type}."
            for head, rel_type, tail in triples
        ]

        return format_str_list_to_a_string(answers, cls.sep)

    def format_relations_for_simple_template(
        self, ralations: Optional[List[Dict[str, List[str]]]]
    ):
        """输出模板为简单模板，即 head (head_type) rel_type tail (tail_type)"""
        # 示例输出： Havana (location) <organization_based_in> Radio Reloj Network (organization)
        if ralations is None or len(ralations) == 0:
            return ""

        relation_output = []

        for r in ralations:
            head = Entity.from_json(r["head"])
            tail = Entity.from_json(r["tail"])

            rel_type = (
                self.label2token.get(r["type"].lower())
                if self.use_special_token_for_label
                else r["type"].lower()
            )
            relation_output.append(
                f"{head.name} ({head.type}) {rel_type} {tail.name} ({tail.type})"
            )

        return format_str_list_to_a_string(relation_output, self.sep)

    def fromat_relations_for_re_task(
        self,
        relations: Optional[List[Dict[str, List[str]]]],
    ) -> Optional[str]:
        """
        为关系抽取任务输出答案
        根据 re_strict 参数，分为 RES 和 RE 两种输出模板
        RES输出模板为：{head} is a {head_type}, {tail} is a {tail_type}, {head} {rel_type} {tail}
            按照输出全部实体，再输出关系的顺序，实体内部输出顺序：head -> tail ,
            关系输出顺序： 三元组出现顺序
        RE输出模板为：{head} {rel_type} {tail}
        """
        if relations is None or len(relations) == 0:
            return ''

        # 是否为严格关系抽取任务，如果是，则需要找出所有的实体及其类型
        entity_output = ""
        if self.template_order in [0, 1]:
            # 找出所有的实体及其类型
            entities = set()
            for r in relations:
                entities.add(Entity.from_json(r["head"]))
                entities.add(Entity.from_json(r["tail"]))

            entity_output = EntityFormatter.format_entities(
                entities, template_order=self.template_order
            )

        relation_output = self.from_relations_to_subject_verb_objext_structure(
            relations
        )

        return entity_output + self.sep + relation_output

    def format_relation_quintuples_for_inside_order(
        self, relations: Optional[List[Dict[str, List[str]]]]
    ):
        """为RES 任务，template_order=1的任务输出答案
        它使用主谓宾结构输出关系，实体输出句子，
        针对三元组顺序输出，三元组内部顺序：头实体，尾实体，关系
        """
        if relations is None or len(relations) == 0:
            return ""

        output4ecah_relation = []

        for rel in relations:

            entities = [
                Entity.from_json(rel["head"]),
                Entity.from_json(rel["tail"]),
            ]

            entity_output = EntityFormatter.format_entities(
                entities, template_order=self.template_order
            )

            relation_output = self.from_relations_to_subject_verb_objext_structure(
                relations=[rel]
            )

            output4ecah_relation.append(f"{entity_output}{self.sep}{relation_output}")

        return format_str_list_to_a_string(output4ecah_relation, self.sep)

    def format_answer_for_re(
        self, relations: Optional[List[Dict[str, List[str]]]]
    ) -> str:
        triples = self.extract_triples(relations=relations)
        answers = [f"{head} {rel_type} {tail}" for head, rel_type, tail in triples]

        return format_str_list_to_a_string(answers, self.sep)


if __name__ == "__main__":
    # print(AnswerFormatter.from_relations({'head': [{'name': 'Lionel Messi'}, {'type': 'footballer'}, {'name': 'Barcelona'}], 'tail': [{'name': 'Lionel Messi'}, {'type': 'footballer'}, {'name': 'Barcelona'}]}, 'en'))

    # Output: 'Lionel Messi is a footballer and Barcelona is a footballer.
    example = [
        {
            "sentence": "In 1863 , American automaker Henry Ford was born in Dearborn Township , Mich .",
            "relations": [
                {
                    "head": {"name": "Henry Ford", "type": "People", "pos": [29, 39]},
                    "type": "Live in",
                    "tail": {
                        "name": "Dearborn Township , Mich",
                        "type": "Location",
                        "pos": [52, 76],
                    },
                }
            ],
        },
    ]
    answer = AnswerFormatter.fromat_relations_to_relation_type(example[0]["relations"])

    print(answer)

    # print(
    #     parse_relation_triple_from_output2(answer, ["live in"]),
    # )
    # # print(SemvalAnswerFormatter.from_relations(example["relations"], "en"))

    # # print(EntityFormatter.from_relations(example["relations"], "en"))
